from typing import TYPE_CHECKING, Optional, Dict

from .builtin_classes import List
from .cython_interface.shared import alloc_obj_memory, alloc_array_memory, incref, int_array_to_list

if TYPE_CHECKING:
    from MelodieAllegro.types import PMObjectType
    from Melodie import AgentList, Agent, Scenario, Model

ObjType = Optional["PMObjectType"]


class O:
    pass


class Converter:
    def __init__(self,
                 scenario_obj_type: ObjType,
                 agent_list_obj_type: ObjType):
        self.scenario_ptr = -1
        self.scenario_obj_type: ObjType = scenario_obj_type
        self.agent_list_obj_type: ObjType = agent_list_obj_type
        self.agent_list_type: ObjType = None
        self.agent_obj_types: Dict[str, ObjType] = {}
        self.list_types: Dict[str, ObjType] = {}

        self.model_props_ptrs: Dict[str, int] = {}

    def convert_agentlist_to_buffer(self, agent_list: "AgentList"):
        pass

    def convert_scenario_to_buffer(self, scenario: "Scenario"):
        return self.scenario_obj_type.create_in_memory(scenario.to_dict())

    def convert_list_of_agent_to_buffer(self, list_of_agent: "List[Agent]",
                                        agent_obj_type: "PMObjectType",
                                        list_obj_type: "PMObjectType",
                                        scenario_ptr: int):
        agent_ptrs = []
        for agent in list_of_agent:
            agent_ptr = self.convert_agent_to_buffer(agent, agent_obj_type, scenario_ptr)
            agent_ptrs.append(agent_ptr)
        arr_id = alloc_array_memory(agent_ptrs + [0 for i in range(len(agent_ptrs))])

        return list_obj_type.create_in_memory(
            {'length': len(agent_ptrs), 'max_length': len(agent_ptrs) * 2, 'data': arr_id})

    def convert_agent_list(self, agent_list: "AgentList", scenario_ptr: int):
        a: Agent = agent_list[0]
        agent_cls_protocol_name = a.__class__.__name__
        agent_cls_name = agent_cls_protocol_name  # [:len(agent_cls_protocol_name) - len('Protocol')]
        agents_ptr = self.convert_list_of_agent_to_buffer(agent_list.agents,
                                                          self.agent_obj_types[agent_cls_name],
                                                          self.list_types[agent_cls_name], scenario_ptr)
        return self.agent_list_obj_type.create_in_memory({'agents': agents_ptr})

    def convert_agent_to_buffer(self, agent: "Agent", agent_obj_type: "PMObjectType",
                                scenario_ptr: int):
        d = {}
        d.update(agent.__dict__)
        d.update({"id": agent.id, "scenario": scenario_ptr})
        incref(scenario_ptr)
        return agent_obj_type.create_in_memory(d)

    def update_agent_from_buffer(self, agent_autoptr: int, agent_obj_type: "PMObjectType", agent: "Agent"):
        agent_obj_type.write_to_python_object(agent_autoptr, agent)

    def update_agent_list_from_buffer(self, agent_list_ptr: int, agent_list_obj_type: "PMObjectType",
                                      agent_list: "AgentList"):
        """
        将AgentList的值从MA中读取到Python空间，并且更新agent_list

        :param agent_list_ptr:
        :param agent_list_obj_type:
        :param agent_list:
        :return:
        """
        obj = O()
        agent_list_obj_type.write_to_python_object(agent_list_ptr, obj, ignore_object=False)
        agents_ptr = obj.__dict__['agents']  # type `List`

        a: Agent = agent_list[0]
        agent_cls_name = a.__class__.__name__

        list_obj = O()
        self.list_types[agent_cls_name].write_to_python_object(agents_ptr, list_obj, False)
        # print('obj_result', list_obj.__dict__)

        length = list_obj.__dict__['length']
        agents_array_ptr = list_obj.__dict__['data']

        agents_ptr_list = int_array_to_list(agents_array_ptr)
        for agent_ptr in agents_ptr_list[:length]:
            agent_data = O()
            self.agent_obj_types[agent_cls_name].write_to_python_object(agent_ptr, agent_data, False)
            dic = agent_data.__dict__
            agent_id = dic.pop('id')
            # for k, v in dic.items():
            agent: Agent = agent_list.get_agent(agent_id)
            agent.set_params(dic)
        # print(agents_ptr_list)

    def update_model_from_buffer(self, model):
        from Melodie import AgentList
        for prop_name, ptr in self.model_props_ptrs.items():
            attr = getattr(model, prop_name)
            prop_type = type(attr)
            if prop_type == AgentList:
                agent_list = attr
                self.update_agent_list_from_buffer(ptr, self.agent_list_obj_type, agent_list)
        pass

    def convert_model(self, model: "Model", proxy_cls: "PMObjectType", scenario_ptr: int):
        from Melodie import AgentList
        props_dict: Dict[str, int] = {}
        for k, v in model.__dict__.items():
            if isinstance(v, AgentList):
                assert proxy_cls.has_attr(k), f"{proxy_cls} does not have attribute `{k}`"
                agent_list_ptr = self.convert_agent_list(v, scenario_ptr)
                props_dict[k] = agent_list_ptr
                self.model_props_ptrs[k] = agent_list_ptr
        return proxy_cls.create_in_memory(props_dict)
